'''
Description: arm controller
Version: V1.0
Author: zw_1520@163.com
Date: 2024-11-19 15:51:41
LastEditors: zw_1520@163.com
LastEditTime: 2025-04-22 08:30:05
Copyright (C) All rights reserved.
'''

import argparse
import datetime
import signal
import sys
import threading
import time

import aimdk.protocol_aimrt_rpc_pb2
import aimdk.protocol_pb2
import aimrt_py
import yaml

global_aimrt_core = None
running = threading.Event()
running.set()

target_positions_list = [
        [0.24, -1.4, 1.5708, 0.8, 0.00, 0.00, 0.00],
        [0.44, -1.1, 1.3708, 1.0, 0.00, 0.00, 0.00],
        [0.64, -0.8, 1.1708, 1.2, 0.00, 0.00, 0.00],
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00]]




def signal_handler(sig, _):
    global global_aimrt_core, running

    if (global_aimrt_core and (sig == signal.SIGINT or sig == signal.SIGTERM)):
        running.clear()
        global_aimrt_core.Shutdown()
        return

    sys.exit(0)


def arm_move(module_logger, mc_action_proxy, mc_motion_proxy):
    global target_positions_list
    # ==========获取当前 Action==========
    req = aimdk.protocol_pb2.McActionRequest()
    status, rsp = mc_action_proxy.GetAction(req)
    if not status.OK():
        aimrt_py.error(module_logger, f"Get action failed. status: {status.ToString()}")
        return
    aimrt_py.info(module_logger, f"Get action success. action: {rsp}")

    VALID_ACTIONS = (
        aimdk.protocol_pb2.McAction.McAction_STAND_DEFAULT,
        aimdk.protocol_pb2.McAction.McAction_STAND_ARM_EXT_JOINT_TRAJ
    )
    if rsp.info.current_action not in VALID_ACTIONS:
        aimrt_py.error(module_logger, f"Current action is not valid. action: {rsp.info.current_action}")
        return

    # ==========切换到站立执行上肢动作 Action==========
    req = aimdk.protocol_pb2.McActionRequest()
    req.command.action = aimdk.protocol_pb2.McAction.McAction_STAND_ARM_EXT_JOINT_TRAJ
    status, rsp = mc_action_proxy.SetAction(req)

    if not status.OK():
        aimrt_py.error(module_logger, f"Set action failed. status: {status.ToString()}")
        return

    aimrt_py.info(module_logger, f"Set action to STAND_ARM_EXT_JOINT_TRAJ success.")

    # ==========发送上肢动作请求==========

    req = aimdk.protocol_pb2.MoveTrajectoryRequest()
    left_joint_path = req.target.left
    left_joint_path.reference_time = 0.0

    positions_list = [
        [0.24, -1.4, 1.5708, 0.8, 0.00, 0.00, 0.00],
        [0.44, -1.1, 1.3708, 1.0, 0.00, 0.00, 0.00],
        [0.64, -0.8, 1.1708, 1.2, 0.00, 0.00, 0.00],
        [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00]
    ]
    time_since_reference = 0.0
    for positions in target_positions_list:
        point = aimdk.protocol_pb2.ScalarTrajectoryPoint()
        point.time_since_reference = time_since_reference
        point.positions.extend(positions)
        left_joint_path.points.append(point)
        time_since_reference += 1.0

    # Call trajectory move service
    status, rsp = mc_motion_proxy.TrajectoryMove(req)
    if not status.OK():
        aimrt_py.error(module_logger, f"Call TrajectoryMove failed, status: {status.ToString()}")
        return

    aimrt_py.info(module_logger, f"TrajectoryMove success rsp: {rsp}")


def arm_control(config_file_path):
    global target_positions_list
    # signal.signal(signal.SIGINT, signal_handler)
    # signal.signal(signal.SIGTERM, signal_handler)

    print("AimRT start.")

    aimrt_core = aimrt_py.Core()

    global global_aimrt_core
    global_aimrt_core = aimrt_core

    # Initialize
    core_options = aimrt_py.CoreOptions()
    core_options.cfg_file_path = config_file_path
    aimrt_core.Initialize(core_options)

    # Create Module
    module_handle = aimrt_core.CreateModule("ArmModule")

    # Register rpc client
    rpc_handle = module_handle.GetRpcHandle()
    mc_action_proxy = aimdk.protocol_aimrt_rpc_pb2.McActionServiceProxy(rpc_handle)
    ret = mc_action_proxy.RegisterClientFunc(rpc_handle)
    assert ret, "Register pb rpc service failed."

    mc_motion_proxy = aimdk.protocol_aimrt_rpc_pb2.McMotionServiceProxy(rpc_handle)
    ret = mc_motion_proxy.RegisterClientFunc(rpc_handle)
    assert ret, "Register pb rpc service failed."

    # Start
    thread = threading.Thread(target=aimrt_core.Start)
    thread.start()

    # Sleep
    time.sleep(1)

    # 行走
    arm_move(module_handle.GetLogger(), mc_action_proxy, mc_motion_proxy)

    # 退出
    aimrt_core.Shutdown()

    while thread.is_alive():
        thread.join(1.0)

    print("AimRT exit.")


    # Sleep
    time.sleep(1)

    # 行走
    arm_move(module_handle.GetLogger(), mc_action_proxy, mc_motion_proxy)

    # 退出
    aimrt_core.Shutdown()

    while thread.is_alive():
        thread.join(1.0)

    print("AimRT exit.")


if __name__ == '__main__':
    main()
